import os
import numpy as np
from datetime import datetime
from PIL import Image, ImageDraw, ImageFont
from ultralytics import YOLO
from ultralytics.utils.plotting import colors

def soft_nms(dets, scores, sigma=0.5, thresh=0.001, method=2):
    """
    参数:
        dets: 检测框坐标 (N,4) [x1,y1,x2,y2]
        scores: 每个框的置信度 (N,)
        sigma: 高斯函数sigma
        thresh: 分数阈值
        method: 0=线性,1=高斯,2=原始NMS

    返回:
        保留的索引
    """
    N = dets.shape[0]
    indexes = np.arange(N)

    for i in range(N):
        max_pos = i
        max_score = scores[i]
        pos = i + 1

        # 找到最大分数框
        while pos < N:
            if scores[pos] > max_score:
                max_score = scores[pos]
                max_pos = pos
            pos += 1

        # 交换位置
        dets[[i, max_pos]] = dets[[max_pos, i]]
        scores[i], scores[max_pos] = scores[max_pos], scores[i]
        indexes[i], indexes[max_pos] = indexes[max_pos], indexes[i]

        # 抑制处理
        pos = i + 1
        while pos < N:
            x1 = max(dets[i, 0], dets[pos, 0])
            y1 = max(dets[i, 1], dets[pos, 1])
            x2 = min(dets[i, 2], dets[pos, 2])
            y2 = min(dets[i, 3], dets[pos, 3])

            w = max(0.0, x2 - x1 + 1)
            h = max(0.0, y2 - y1 + 1)
            inter = w * h
            iou = inter / ((dets[i, 2] - dets[i, 0]) * (dets[i, 3] - dets[i, 1]) +
                           (dets[pos, 2] - dets[pos, 0]) * (dets[pos, 3] - dets[pos, 1]) - inter + 1e-16)

            if method == 1:  # 高斯
                weight = np.exp(-(iou * iou) / sigma)
            elif method == 0:  # 线性
                if iou > sigma:
                    weight = 1 - iou
                else:
                    weight = 1
            else:  # 原始NMS
                if iou > sigma:
                    weight = 0
                else:
                    weight = 1

            scores[pos] = weight * scores[pos]

            # 如果分数太低则丢弃
            if scores[pos] < thresh:
                dets[[pos, N - 1]] = dets[[N - 1, pos]]
                scores[pos], scores[N - 1] = scores[N - 1], scores[pos]
                indexes[pos], indexes[N - 1] = indexes[N - 1], indexes[pos]
                N -= 1
                pos -= 1
            pos += 1

    return indexes[:N]


def picture_flag(model, picture_path, result_image_dir):
    os.makedirs(result_image_dir, exist_ok=True)

    # 1. 用PIL读取图片
    img_rgb = Image.open(picture_path).convert("RGB")
    img_width, img_height = img_rgb.size

    # 2. 转换为NumPy数组供YOLO处理
    img_np = np.array(img_rgb)  # Shape: (H, W, 3) RGB

    # 3. 使用模型预测（禁用内置NMS）
    results = model.predict(img_np, nms=False)  # 输入RGB格式

    # 4. 提取预测结果
    boxes = results[0].boxes
    dets = boxes.xyxy.cpu().numpy()
    scores = boxes.conf.cpu().numpy()
    cls = boxes.cls.cpu().numpy()

    # 5. 应用Soft-NMS
    if len(dets) > 0:
        keep = soft_nms(dets, scores, sigma=0.5, thresh=0.001, method=1)
        dets = dets[keep]
        scores = scores[keep]
        cls = cls[keep]

    # 6. 在RGB图像上绘制结果（直接用PIL）
    draw = ImageDraw.Draw(img_rgb)
    try:
        font = ImageFont.truetype("arial.ttf", size=20)  # 尝试加载字体
    except:
        font = ImageFont.load_default()  # 回退到默认字体

    for box, score, class_id in zip(dets, scores, cls):
        x1, y1, x2, y2 = map(int, box)
        class_id_int = int(class_id)

        # 获取YOLO默认颜色（已经是RGB格式）
        color_rgb = colors(class_id_int)

        # 绘制矩形框
        draw.rectangle([x1, y1, x2, y2], outline=color_rgb, width=2)

        # 绘制标签
        label = f"{results[0].names[class_id_int]} {score:.2f}"
        draw.text((x1, y1 - 20), label, fill=color_rgb, font=font)

    # 7. 保存RGB格式图片（用PIL）
    current_time = datetime.now().strftime("%m%d%H%M%S")
    output_path = os.path.join(result_image_dir, f"{current_time}_0.jpg")
    img_rgb.save(output_path, quality=95)  # 保持RGB格式保存

    # 8. 返回RGB图像（PIL Image对象或NumPy数组均可）
    return np.array(img_rgb)  # 返回NumPy数组（RGB）


def main():
    # 加载训练好的模型
    model = YOLO("best_cars_detection.pt")

    # 图片路径
    image_path = "img_1.png"

    # 结果保存目录
    result_dir = os.path.join("picture_inference", "picture_flag")

    # 进行检测
    result_image = picture_flag(model, image_path, result_dir)

    print(f"检测完成，结果已保存到: {result_dir}")


if __name__ == "__main__":
    main()